import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader

inputs = torch.tensor([1, 2, 3], dtype=torch.float32)
targets = torch.tensor([1, 2, 5], dtype=torch.float32)

inputs = torch.reshape(inputs, [1, 1, 1, 3])
targets = torch.reshape(targets, [1, 1, 1, 3])

loss = nn.L1Loss()
print(loss(inputs, targets))

train_data = torchvision.datasets.CIFAR10("./dataset", train=True, download=False,
                                          transform=torchvision.transforms.ToTensor())
train_dataloader = DataLoader(train_data, batch_size=1, shuffle=True, num_workers=0, drop_last=False)


class MyMod(nn.Module):
    def __init__(self):
        super(MyMod, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(in_features=1024, out_features=64),
            nn.Linear(in_features=64, out_features=10)
        )

    def forward(self, x):
        output = self.model(x)
        return output

# 这里要注意pytorch计算交叉熵误差的函数自带softmax，训练时模型里面不要添加softmax
# 交叉熵使用的是ln
loss = nn.CrossEntropyLoss()
mymod = MyMod()
for data in train_dataloader:
    imgs, targets = data
    output = mymod(imgs)
    print(output)
    print(targets)
    loss_result = loss(output, targets)
    print(loss_result)
    # 反向传播
    loss_result.backward()
    break
